import torch
import torch.nn as nn


class LeNetModel(nn.Module):
    def __init__(self, num_class=10):
        super(LeNetModel, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, kernel_size=5)
        # 输入通道1，输出通道6
        self.conv2 = nn.Conv2d(6, 16, kernel_size=5)

        self.fc1 = nn.Linear(16 * 4 * 4, 120)
        # 全连接层

        self.fc2 = nn.Linear(120, 84)

        # 输出10类（0-9）
        self.fc3 = nn.Linear(84, num_class)

        self.pool = nn.MaxPool2d(2, 2)  # 池化层

    def forward(self, x):
        # [B, 1, 28, 28]

        x = torch.relu(self.conv1(x))
        # [B, 6, 24, 24]

        x = self.pool(x)
        # [B, 6, 12, 12]

        x = torch.relu(self.conv2(x))
        # [B, 16, 8, 8]

        x = self.pool(x)
        # [B, 16, 4, 4]


        x = x.view(-1, 16 * 4 * 4)  # 展平
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x
